import numpy as np
import matplotlib.pyplot as plt
from scipy.signal import find_peaks

def _create_signal():
    """
    创建一个包含多个波峰的示例信号。
    """
    x = np.linspace(0, 100, 500)
    # 包含不同高度、宽度和突出度的波峰
    y = (np.sin(x / 5) * 2 + np.cos(x / 2) * 0.5 + 
         np.sin(x / 10) * 1.5 + np.random.rand(len(x)) * 0.2)
    
    # 添加一个平顶波峰
    y[200:210] = 2.5 # 模拟平顶
    y[205] = 2.6 # 确保中间有一个峰值
    
    return x, y

def _plot_peaks(x, y, peaks, title, properties=None):
    """
    绘制信号和检测到的波峰。
    """
    plt.figure(figsize=(12, 6))
    plt.plot(x, y, label='原始信号')
    plt.plot(x[peaks], y[peaks], "x", label='检测到的波峰', color='red')
    
    if properties:
        # 绘制波峰的属性，例如突出度基线或宽度线
        if "prominences" in properties:
            for i, peak in enumerate(peaks):
                plt.vlines(x[peak], y[peak] - properties["prominences"][i], 
                           y[peak], color='green', linestyle='--', 
                           label='突出度' if i == 0 else "")
        if "widths" in properties and "width_heights" in properties and "left_ips" in properties and "right_ips" in properties:
            for i, peak in enumerate(peaks):
                plt.hlines(properties["width_heights"][i], 
                           x[int(properties["left_ips"][i])], 
                           x[int(properties["right_ips"][i])], 
                           color='purple', linestyle='-', 
                           label='宽度' if i == 0 else "")

    plt.title(title)
    plt.xlabel('样本点')
    plt.ylabel('幅值')
    plt.legend()
    plt.grid(True)
    plt.rcParams['font.sans-serif'] = ['Microsoft YaHei']  # 更好的中文字体，可以显示“-”号
    plt.show()

def demo_find_peaks_parameters():
    """
    演示scipy.signal.find_peaks不同参数的实际效果。
    """
    x, y = _create_signal()

    # 1. height - 波峰高度条件
    # 只保留高度≥1.5的波峰
    peaks_height_1 = find_peaks(y, height=1.5)[0]
    _plot_peaks(x, y, peaks_height_1, 'find_peaks: height=1.5')

    # 保留高度在0.5到2.0之间的波峰
    peaks_height_2 = find_peaks(y, height=(0.5, 2.0))[0]
    _plot_peaks(x, y, peaks_height_2, 'find_peaks: height=(0.5, 2.0)')

    # 2. threshold - 阈值条件
    # 波峰必须比相邻点高出至少0.5
    peaks_threshold = find_peaks(y, threshold=0.2)[0]
    _plot_peaks(x, y, peaks_threshold, 'find_peaks: threshold=0.5')

    # 3. distance - 距离条件
    # 相邻波峰至少间隔50个样本点
    peaks_distance = find_peaks(y, distance=50)[0]
    _plot_peaks(x, y, peaks_distance, 'find_peaks: distance=50')

    # 4. prominence - 突出度条件
    # 只保留突出度≥1.0的波峰
    peaks_prominence_1, properties_prominence_1 = find_peaks(y, prominence=1.0)
    _plot_peaks(x, y, peaks_prominence_1, 'find_peaks: prominence=1.0', properties=properties_prominence_1)

    # 保留突出度在0.5到1.5之间的波峰
    peaks_prominence_2, properties_prominence_2 = find_peaks(y, prominence=(0.5, 1.5))
    _plot_peaks(x, y, peaks_prominence_2, 'find_peaks: prominence=(0.5, 1.5)', properties=properties_prominence_2)

    # 5. width - 宽度条件
    # 保留宽度在10到50个样本之间的波峰
    peaks_width, properties_width = find_peaks(y, width=(10, 50))
    _plot_peaks(x, y, peaks_width, 'find_peaks: width=(10, 50)', properties=properties_width)

    # 6. plateau_size - 平台大小
    # 平台至少5个样本宽
    peaks_plateau = find_peaks(y, plateau_size=5)[0]
    _plot_peaks(x, y, peaks_plateau, 'find_peaks: plateau_size=5')

    # 7. wlen - 窗口长度 (影响 prominence 和 width 的计算)
    # 较小的wlen可能导致更多局部峰被认为是突出峰
    peaks_wlen_small, properties_wlen_small = find_peaks(y, prominence=0.5, wlen=50)
    _plot_peaks(x, y, peaks_wlen_small, 'find_peaks: prominence=0.5, wlen=50', properties=properties_wlen_small)

    # 较大的wlen可能平滑掉一些局部细节
    peaks_wlen_large, properties_wlen_large = find_peaks(y, prominence=0.5, wlen=200)
    _plot_peaks(x, y, peaks_wlen_large, 'find_peaks: prominence=0.5, wlen=200', properties=properties_wlen_large)

    # 8. rel_height - 相对高度 (影响 width 的计算)
    # 在0.2相对高度处测量宽度
    peaks_rel_height_02, properties_rel_height_02 = find_peaks(y, prominence=0.5, rel_height=0.2)
    _plot_peaks(x, y, peaks_rel_height_02, 'find_peaks: prominence=0.5, rel_height=0.2', properties=properties_rel_height_02)

    # 在0.8相对高度处测量宽度
    peaks_rel_height_08, properties_rel_height_08 = find_peaks(y, prominence=0.5, rel_height=0.8)
    _plot_peaks(x, y, peaks_rel_height_08, 'find_peaks: prominence=0.5, rel_height=0.8', properties=properties_rel_height_08)

    # 组合参数示例
    peaks_combined, properties_combined = find_peaks(y, height=(0.5, None), distance=20, prominence=0.8, width=(5, 50))
    _plot_peaks(x, y, peaks_combined, 'find_peaks: 组合参数示例', properties=properties_combined)

if __name__ == '__main__':
    demo_find_peaks_parameters()